【妙用协程】 - 可resume的异常处理

标准的异常处理是这样的

try:
    print('hello')
    raise Exception()
    print('!!!')
except:
    print('world')
    print('???')

这段代码会打印出???而不会打印出!!!，因为异常会中断当前流程，跳转到except部分去继续执行。但是有的时候我们希望的这样的行为：

try:
    print('hello')
    print(Scheduler.interrupt())
    print('!!!')
except ProcessInterrupt as pi:
    pi.resume('world')
    print('???')

这段代码打印出!!!而不是???，因为resume的时候把执行重新跳转回interrupt的地方了。这种行为类似vba里的on error resume next（https://msdn.microsoft.com/en-us/library/5hsw66as.aspx）。

如何实现的？其实原理上很简单。interrupt的时候把当前协程的状态保存起来（pickle.dumps），如果决定要resume，就把协程interrupt的时刻的状态重新恢复（pickle.loads）然后从那个点继续执行。

完整的代码（需要pypy或者stackless python）：

import greenlet
import cPickle as pickle
import traceback
import threading
import functools


class ProcessInterrupt(Exception):
    def __init__(self, interruption_point, pi_args):
        self.interruption_point = interruption_point
        self.stacktrace = traceback.extract_stack()
        self.pi_args = pi_args

    def resume(self, resume_with=None):
        Scheduler.resume(self.interruption_point, resume_with)

    def __repr__(self):
        return '>>>ProcessInterrupt>>>%s' % repr(self.stacktrace)

    def __str__(self):
        return repr(self)

    def __unicode__(self):
        return repr(self)


class Scheduler(object):
    current = threading.local()

    def __init__(self):

        if getattr(self.current, 'instance', None):
            raise Exception('can not have two scheduler in one thread')
        self.scheduler_greenlet = greenlet.getcurrent()
        self.current.instance = self

    def __call__(self, action, action_args):
        next = action, action_args
        while next:
            action, action_args = next
            if 'init' == action:
                next = action_args['init_greenlet'].switch()
            elif 'interrupt' == action:
                interruption_point = pickle.dumps(action_args['switched_from'])
                should_resume, resume_with = False, None
                next = action_args['switched_from'].switch(
                    should_resume, resume_with, interruption_point)
            elif 'resume' == action:
                should_resume, resume_with, interruption_point = True, action_args['resume_with'], action_args[
                    'interruption_point']
                next = pickle.loads(action_args['interruption_point']).switch(
                    should_resume, resume_with, interruption_point)
            else:
                raise NotImplementedError('unknown action: %s' % action)

    @classmethod
    def start(cls, init_func, *args, **kwargs):
        scheduler = Scheduler()
        init_greenlet = greenlet.greenlet(functools.partial(init_func, *args, **kwargs))
        scheduler('init', {
            'init_greenlet': init_greenlet,
        })

    @classmethod
    def interrupt(cls, pi_args=None):
        should_resume, resume_with, interruption_point = cls.switch_to_scheduler('interrupt', {
            'switched_from': greenlet.getcurrent()
        })
        if should_resume:
            return resume_with
        else:
            pi = ProcessInterrupt(interruption_point, pi_args)
            raise pi

    @classmethod
    def resume(cls, interruption_point, resume_with=None):
        cls.switch_to_scheduler('resume', {
            'interruption_point': interruption_point,
            'resume_with': resume_with
        })

    @classmethod
    def switch_to_scheduler(cls, *args, **kwargs):
        return cls.current.instance.scheduler_greenlet.switch(*args, **kwargs)


if '__main__' == __name__:
    def init():
        try:
            print('hello')
            print(Scheduler.interrupt())
            print('!!!')
        except ProcessInterrupt as pi:
            pi.resume('world')
            print('???')

        try:
            print('hello')
            raise Exception()
            print('!!!')
        except:
            print('world')
            print('???')

    Scheduler.start(init)